#!/usr/bin/env python
# coding: utf-8

# In[1]:


#from sympy import*
import math
import numpy as np
from functools import reduce
np.set_printoptions(precision=4, suppress=True)


# In[2]:


#正向运动学计算（包含求解雅克比矩阵）
def fk(joints): #输入参数：关节角joints
                #输出参数：变换矩阵T_4，雅克比矩阵J
    theta_1, theta_2, theta_3, theta_4= joints
# DH_pramater: [theta, a, d, alpha]，注意这里的单位是m
    DH = [[theta_1, 0.008, 0, np.pi/2],
          [theta_2, 0.0342, 0.232, np.pi/2],
          [theta_3, 0.385,0,0],
          [theta_4, 0.082, 0, -np.pi/2]]
#根据DH参数计算变换矩阵
    T=[]
    T_=[]
    for i in range(4):
        theta,d,a,alpha=DH[i]
        cth,sth=np.cos(theta),np.sin(theta)
        cal,sal=np.cos(alpha),np.sin(alpha)
        t=np.array([[cth, -sth*cal, sth*sal, cth*a],
                     [sth, cth*cal, -cth*sal, sth*a],
                     [0, sal, cal, d],
                     [0, 0, 0, 1]])
        T.append(t)
#根据DH参数计算变换矩阵对各关节角的偏导数
        t_=np.array([[-sth, -cth*cal, cth*sal, -sth*a],
                      [cth, -sth*cal, sth*sal, cth*a],
                     [0, 0, 0, 0],
                     [0, 0, 0, 0]])
        T_.append(t_)
    T_4 =reduce(np.dot,T)
#避免使用python自带的求导函数（速度慢），此处直接给出变换矩阵对各关节角求导后的矩阵形式
    J_1=T_[0].dot(T[1]).dot(T[2]).dot(T[3])
    J_2=T[0].dot(T_[1]).dot(T[2]).dot(T[3])
    J_3=T[0].dot(T[1]).dot(T_[2]).dot(T[3])
    J_4=T[0].dot(T[1]).dot(T[2]).dot(T_[3])
    J_=[J_1, J_2, J_3, J_4]
#由旋转矩阵得到机械臂末端6个自由度关于关节角的函数表达式
    p_x=T_4[0,3]
    p_y=T_4[1,3]
    p_z=T_4[2,3] #三维坐标位置
    alpha=math.atan2(-1*T_4[1,2],T_4[2,2])
    beta=math.asin(T_4[0,2])
    gama=math.atan2(-1*T_4[0,1],T_4[0,0]) #绕三个坐标轴的转角（欧拉角）
    f=np.array([p_x,p_y,p_z,alpha,beta,gama])
#构造速度雅克比矩阵
    J=np.zeros([6,4])
    for i in range(4):
        J[0:3,i]=J_[i][0:3,3]
        J[3,i]=(-1*J_[i][1,2]*T_4[2,2]+J_[i][2,2]*T_4[1,2])/(T_4[1,2]**2+T_4[2,2]**2)
        J[4,i]=J_[i][0,2]/(1-T_4[0,2]**2)**0.5
        J[5,i]=(-1*J_[i][0,1]*T_4[0,0]+J_[i][0,0]*T_4[0,1])/(T_4[0,1]**2+T_4[0,0]**2)
    return (f,J) 


# In[3]:


#逆运动学计算      
def ik(f_tar, joints_init, tolerance = 1e-7): #输入参数：目标位姿的变换矩阵T_tar，当前机械臂的关节角joints_init，默认容许误差
                                              #输出参数：达到目标位姿所需的关节角joints，迭代次数itertime                                    
    itertime = 0
    joints = joints_init.copy()
    while itertime < 100:
        (f_cur,jac) = fk(joints)
        deltaf =f_tar - f_cur
        error = deltaf.dot(np.transpose(deltaf))
        if error < tolerance:
            return (joints,itertime)
        deltaq = np.linalg.pinv(jac).dot(deltaf) #求雅克比矩阵伪逆，计算迭代公式中的关节角变化量
        joints = joints + deltaq
        itertime += 1
    return False


# In[4]:


#测试计算结果
joints_tar = np.array([0, -np.pi/2, 0, 0])
(f_tar,J_tar) = fk(joints_tar)
joints_init = joints_tar + np.random.uniform(-0.3, 0.3, 4) #设置初始关节角度位于期望关节角度附近的领域内
(q,time) = ik(f_tar, joints_init)
get_ipython().run_line_magic('timeit', 'ik(f_tar, joints_init)')
print(q)


# In[ ]:




